package com.orhanobut.hawk.encryption

import java.nio.charset.StandardCharsets
import java.security.MessageDigest
import java.security.SecureRandom
import javax.crypto.Cipher
import javax.crypto.spec.GCMParameterSpec
import javax.crypto.spec.SecretKeySpec

class AesEncryption(keyBytes: ByteArray) : Encryption {
  private val keySpec: SecretKeySpec = SecretKeySpec(MessageDigest.getInstance("SHA-256").digest(keyBytes), "AES")
  private val random = SecureRandom()

  override fun init(): Boolean = true

  override fun encrypt(key: String, value: String): String {
    val iv = ByteArray(12).also { random.nextBytes(it) }
    val cipher = Cipher.getInstance("AES/GCM/NoPadding")
    cipher.init(Cipher.ENCRYPT_MODE, keySpec, GCMParameterSpec(128, iv))
    cipher.updateAAD(key.toByteArray(StandardCharsets.UTF_8))
    val ct = cipher.doFinal(value.toByteArray(StandardCharsets.UTF_8))
    return toHex(iv) + ":" + toHex(ct)
  }

  override fun decrypt(key: String, value: String): String {
    val parts = value.split(":")
    val iv = fromHex(parts[0])
    val ct = fromHex(parts[1])
    val cipher = Cipher.getInstance("AES/GCM/NoPadding")
    cipher.init(Cipher.DECRYPT_MODE, keySpec, GCMParameterSpec(128, iv))
    cipher.updateAAD(key.toByteArray(StandardCharsets.UTF_8))
    val pt = cipher.doFinal(ct)
    return String(pt, StandardCharsets.UTF_8)
  }

  private fun toHex(bytes: ByteArray): String {
    val sb = StringBuilder(bytes.size * 2)
    for (b in bytes) {
      val v = b.toInt() and 0xFF
      if (v < 16) sb.append('0')
      sb.append(Integer.toHexString(v))
    }
    return sb.toString()
  }

  private fun fromHex(hex: String): ByteArray {
    val len = hex.length
    val out = ByteArray(len / 2)
    var i = 0
    var j = 0
    while (i < len) {
      out[j++] = ((hexDigit(hex[i]) shl 4) or hexDigit(hex[i + 1])).toByte()
      i += 2
    }
    return out
  }

  private fun hexDigit(c: Char): Int {
    return when (c) {
      in '0'..'9' -> c.code - '0'.code
      in 'a'..'f' -> c.code - 'a'.code + 10
      in 'A'..'F' -> c.code - 'A'.code + 10
      else -> 0
    }
  }
}